#include <algorithm>

#include "TrackingAlgorithm.h"
#include "networkfunc.h"
#include "TrackerManager.h"
#include "nn_matching.h"

using namespace std;

struct TrackingAlgorithm::TrackData {
    KalmanTracker kalman;
    FeatureBundle feats;
};

TrackingAlgorithm::TrackingAlgorithm(int width, int height, void *detector, void *extractor):
    height(height), width(width), obj_det(detector), re_id(extractor),feat_metric(make_unique<FeatureMetric<TrackData>>(data)) {
    array<int64_t, 2> ori_dim{int64_t(height), int64_t(width)};
    manager= std::make_unique<TrackerManager<TrackData>>(data, ori_dim);
}

TrackingAlgorithm::~TrackingAlgorithm() = default;

std::vector<std::vector<cv::Rect2f>> TrackingAlgorithm::getDetections(std::vector<std::vector<Detection>> &pedestrianbox)
{
    std::vector<std::vector<cv::Rect2f>> dets;
    for (int b = 0; b < pedestrianbox.size(); b++) {
        std::vector<cv::Rect2f> box_per_frame;
        for (size_t i = 0; i < pedestrianbox[b].size(); i++) {
            Detection body = pedestrianbox[b][i];
            float x1 = body.box.x;
            float y1 = body.box.y;
            float x2 = x1 + body.box.width;
            float y2 = y1 + body.box.height;
            if (body.cls == 1) {//body
                box_per_frame.push_back(body.box);
            }
        }
        dets.push_back(box_per_frame);
    }

    return dets;
}

static float shiftIOU(const cv::Rect2f &bb_test, const cv::Rect2f &bb_gt) {
    auto in = (bb_test & bb_gt).area();
    auto un = bb_test.area() + bb_gt.area() - in;

    if (un < DBL_EPSILON)
        return 0;

    return in / bb_test.area();
}

std::vector<Detection> TrackingAlgorithm::update(const cv::Mat &ori_img, std::vector<std::vector<Detection>> &pedestrianbox){
    manager->predict();
    manager->remove_nan();
    NetworkFeatExtract *extractor = (NetworkFeatExtract*)re_id;
    NetworkDetect *detector = (NetworkDetect *)obj_det;
    std::vector<cv::Mat> imgs{ori_img};
//    double start = static_cast<double>(cv::getTickCount());
    pedestrianbox.clear();
    detector->detect(imgs, pedestrianbox, 512, 512);
//    double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
//    cout << "time: " << time << "ms" << endl;

    //getDetections
    std::vector<std::vector<cv::Rect2f>> detections = getDetections(pedestrianbox);

    auto matched = manager->update(
            detections[0],
            [this, &ori_img, &extractor, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) {
                vector<cv::Rect2f> trks;
                for (auto t : trk_ids) {
                    trks.push_back(data[t].kalman.rect());
                }

                vector<cv::Mat> boxes;
                vector<cv::Rect2f> dets;
                for (auto d:det_ids) {
                    dets.push_back(detections[0][d]);
                    cv::Mat tmp;
                    if (d > detections[0].size()-1) {
                        continue;
                    }
                    cv::Rect rect = detections[0][d];
                    cv::resize(ori_img(rect), tmp, cv::Size(128, 384));
                    boxes.push_back(tmp);
                }

                auto iou_mat = iou_dist(dets, trks);
                std::vector<std::vector<float>> feats = extractor->extract(boxes);
                auto feat_mat = feat_metric->distance(feats, trk_ids);
                for (size_t i=0; i<feat_mat.size(); i++) {
                    for (size_t j = 0; j < feat_mat[i].size(); j++) {
                        if ((iou_mat[i][j] > 0.8f) || (feat_mat[i][j] > 0.2f))
                            feat_mat[i][j] = INVALID_DIST;
                    }
                }
                return feat_mat;
            },
            [this, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) {
                vector<cv::Rect2f> trks;
                for (auto t : trk_ids) {
                    trks.push_back(data[t].kalman.rect());
                }
                vector<cv::Rect2f> dets;
                for (auto &d:det_ids) {
                    dets.push_back(detections[0][d]);
                }
                std::vector<std::vector<float>> iou_mat = iou_dist(dets, trks);
                for (size_t i=0; i<iou_mat.size(); i++) {
                    for (size_t j = 0; j < iou_mat[i].size(); j++) {
                        if (iou_mat[i][j] > 0.7f)
                            iou_mat[i][j] = INVALID_DIST;
                    }
                }
                //iou_mat.masked_fill_(iou_mat > 0.7f, INVALID_DIST);
                return iou_mat;
            });

    vector<cv::Mat> boxes;
    vector<int> targets;
    for (auto item : matched) {
        targets.emplace_back(std::get<0>(item));
        int index = std::get<1>(item);
        if (index > detections[0].size()-1) {
            continue;
        }
        cv::Mat crop = ori_img(detections[0][index]);
        cv::resize(crop, crop, cv::Size(128, 384));
        boxes.emplace_back(crop);
    }
    std::vector<std::vector<float>> feats = extractor->extract(boxes);
    feat_metric->update(feats, targets);

    manager->remove_deleted();

    std::vector<Detection> tracking = manager->visible_tracks(pedestrianbox);
    for (auto &track : tracking) {
        for (auto &assistant : tracking_assistant){
            if (track.tracking_id == assistant.tracking_id){
                track.state = assistant.state;
                break;
            }
        }
    }

    std::vector<Detection> arm_left, arm_right;

    for (auto det : pedestrianbox[0]){
        if (det.cls == 3) arm_left.push_back(det);
        if (det.cls == 5) arm_right.push_back(det);
    }

    for (auto &det : tracking) {
        for (auto left : arm_left) {
            float iou_ = shiftIOU(left.box, det.box);
            if (iou_ > 0.7) {
                det.state = 1;
                break;
            }
        }
        for (auto right : arm_right) {
            float iou_ = shiftIOU(right.box, det.box);
            if (iou_ > 0.7) {
                if (det.state == 1) {
                    det.state = 2;
                    break;
                }
                else {
                    det.state = 0;
                    break;
                }
            }
        }
    }
    tracking_assistant = tracking;

    return tracking;
}

vector<vector<float>> TrackingAlgorithm::update(const cv::Mat &ori_img) {
    manager->predict();
    manager->remove_nan();
    NetworkFeatExtract *extractor = (NetworkFeatExtract*)re_id;
    NetworkDetect *detector = (NetworkDetect *)obj_det;
    std::vector<std::vector<Detection>> pedestrianbox;
    std::vector<cv::Mat> imgs{ori_img};
//    double start = static_cast<double>(cv::getTickCount());
    detector->detect(imgs, pedestrianbox, 512, 512);
//    double time = ((double)cv::getTickCount() - start) / cv::getTickFrequency();
//    cout << "time: " << time << "ms" << endl;

    //getDetections
    std::vector<std::vector<cv::Rect2f>> detections = getDetections(pedestrianbox);

    auto matched = manager->update(
            detections[0],
            [this, &ori_img, &extractor, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) {
                vector<cv::Rect2f> trks;
                for (auto t : trk_ids) {
                    trks.push_back(data[t].kalman.rect());
                }
                
                vector<cv::Mat> boxes;
                vector<cv::Rect2f> dets;
                for (auto d:det_ids) {
                    dets.push_back(detections[0][d]);
                    cv::Mat tmp;
                    if (d > detections[0].size()-1) {
                        std::cout << " %%%%%%%% " << d << std::endl;
                        continue;
                    }
                    cv::Rect rect = detections[0][d];
                    cv::resize(ori_img(rect), tmp, cv::Size(128, 384));
                    boxes.push_back(tmp);
                }

                auto iou_mat = iou_dist(dets, trks);
                std::vector<std::vector<float>> feats = extractor->extract(boxes);
                auto feat_mat = feat_metric->distance(feats, trk_ids);
//                if (feat_mat.size() > 0)
//                    std::cout << "distance done " << feat_mat.size() << " " << feat_mat[0].size() << " " << det_ids.size() << " " << feats.size() << std::endl;
                for (size_t i=0; i<feat_mat.size(); i++) {
                    for (size_t j = 0; j < feat_mat[i].size(); j++) {
                        if ((iou_mat[i][j] > 0.8f) || (feat_mat[i][j] > 0.2f))
                            feat_mat[i][j] = INVALID_DIST;
                    }
                }
//        feat_mat.masked_fill_((iou_mat > 0.8f).__ior__(feat_mat > 0.2f), INVALID_DIST);
                return feat_mat;
            },
            [this, &detections](const std::vector<int> &trk_ids, const std::vector<int> &det_ids) {
                vector<cv::Rect2f> trks;
                for (auto t : trk_ids) {
                    trks.push_back(data[t].kalman.rect());
                }
                vector<cv::Rect2f> dets;
                for (auto &d:det_ids) {
                    dets.push_back(detections[0][d]);
                }
                std::vector<std::vector<float>> iou_mat = iou_dist(dets, trks);
                for (size_t i=0; i<iou_mat.size(); i++) {
                    for (size_t j = 0; j < iou_mat[i].size(); j++) {
                        if (iou_mat[i][j] > 0.7f)
                            iou_mat[i][j] = INVALID_DIST;
                    }
                }
                //iou_mat.masked_fill_(iou_mat > 0.7f, INVALID_DIST);
                return iou_mat;
            });

    vector<cv::Mat> boxes;
    vector<int> targets;
    for (auto item : matched) {
        targets.emplace_back(std::get<0>(item));
        int index = std::get<1>(item);
        if (index > detections[0].size()-1) {
            continue;
        }
        cv::Mat crop = ori_img(detections[0][index]);
        cv::resize(crop, crop, cv::Size(128, 384));
        boxes.emplace_back(crop);
    }
    std::vector<std::vector<float>> feats = extractor->extract(boxes);
    feat_metric->update(feats, targets);

    manager->remove_deleted();

    return manager->visible_tracks();
}

